import keras.backend as K
from keras import Input, Model
from keras.callbacks import EarlyStopping
from keras.layers import Layer, Activation, Embedding, Bidirectional, LSTM, Dropout, Dense, Flatten
import numpy as np

#squash压缩函数和原文不一样，可自己定义



def squash(x, axis=-1):
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    scale = K.sqrt(s_squared_norm + K.epsilon())
    return x / scale

class Capsule(Layer):
    def __init__(self, num_capsule, dim_capsule, routings=3, kernel_size=(9, 1), share_weights=True,
                 activation='default', **kwargs):
        super(Capsule, self).__init__(**kwargs)
        self.num_capsule = num_capsule
        self.dim_capsule = dim_capsule
        self.routings = routings
        self.kernel_size = kernel_size
        self.share_weights = share_weights
        if activation == 'default':
            self.activation = squash
        else:
            self.activation = Activation(activation)

    def build(self, input_shape):
        super(Capsule, self).build(input_shape)
        input_dim_capsule = input_shape[-1]
        if self.share_weights:
            self.W = self.add_weight(name='capsule_kernel',
                                     shape=(1, input_dim_capsule,
                                            self.num_capsule * self.dim_capsule),
                                     # shape=self.kernel_size,
                                     initializer='glorot_uniform',
                                     trainable=True)
        else:
            input_num_capsule = input_shape[-2]
            self.W = self.add_weight(name='capsule_kernel',
                                     shape=(input_num_capsule,
                                            input_dim_capsule,
                                            self.num_capsule * self.dim_capsule),
                                     initializer='glorot_uniform',
                                     trainable=True)

    def call(self, u_vecs):
        if self.share_weights:
            u_hat_vecs = K.conv1d(u_vecs, self.W)
        else:
            u_hat_vecs = K.local_conv1d(u_vecs, self.W, [1], [1])

        batch_size = K.shape(u_vecs)[0]
        input_num_capsule = K.shape(u_vecs)[1]
        u_hat_vecs = K.reshape(u_hat_vecs, (batch_size, input_num_capsule,
                                            self.num_capsule, self.dim_capsule))
        u_hat_vecs = K.permute_dimensions(u_hat_vecs, (0, 2, 1, 3))

        b = K.zeros_like(u_hat_vecs[:, :, :, 0])  # shape = [None, num_capsule, input_num_capsule]
        #动态路由部分
        for i in range(self.routings):
            b = K.permute_dimensions(b, (0, 2, 1))  # shape = [None, input_num_capsule, num_capsule]
            c = K.softmax(b)
            c = K.permute_dimensions(c, (0, 2, 1))
            b = K.permute_dimensions(b, (0, 2, 1))
            outputs = self.activation(K.batch_dot(c, u_hat_vecs, [2, 2]))
            if i < self.routings - 1:
                b = K.batch_dot(outputs, u_hat_vecs, [2, 3])

        return outputs

    def compute_output_shape(self, input_shape):
        return (None, self.num_capsule, self.dim_capsule)

class LstmCap:
    def __init__(self,vocab_size=5000,emb_dim=300,max_len=600,n_cap=100,cap_dim=100,n_class=10):
        self.vocab_size = vocab_size
        self.emb_dim = emb_dim
        self.max_len = max_len
        self.n_cap = n_cap
        self.cap_dim = cap_dim
        self.n_class = n_class

    def build_model(self):
        word_input = Input(shape=(self.max_len,), dtype="int32")
        embed = Embedding(input_dim=self.vocab_size,
                          output_dim=100,
                          input_length=self.max_len
                          )(word_input)
        x = Bidirectional(LSTM(100, return_sequences=True))(embed)
        x = Capsule(
            num_capsule=self.n_cap, dim_capsule=self.cap_dim,
            routings=3, share_weights=True)(x)
        x = Flatten()(x)
        x = Dropout(0.5)(x)
        outputs = Dense(self.n_class, activation='softmax')(x)
        model = Model(inputs=word_input, outputs=outputs)
        model.compile(loss='categorical_crossentropy', optimizer='nadam', metrics=['accuracy'])
        model.summary()
        return model

    def train(self,dataset,batch_size=128,epochs=10):
        train_seq_mat, train_y, val_seq_mat,val_y,test_seq_mat,test_y =dataset
        early_stop = EarlyStopping(monitor='val_loss',min_delta=0.0001)

        model = self.build_model()
        model_fit =model.fit(
            x=train_seq_mat,
            y=train_y,
            verbose=2,
            batch_size=128,
            epochs=10,
            validation_data=(val_seq_mat,val_y),
            callbacks=[early_stop]
        )
        self.test(model,test_seq_mat,test_y)
        model.save('my_lstm_cap.h5')
if __name__ == '__main__':
    lstm_cap = LstmCap({'1':1},100,200,100,100,3)
    lstm_cap.build_model()

